# coding: UTF-8
import os
import sys
import time
import httpx
import shutil
import tarfile
import argparse
from tqdm import tqdm
import asyncio
import datetime


class MultiCoroutineDownloader(object):
    def __init__(self, url, file_path, cor_num = 20):
        '''
        @param url: 下载链接
        @param file_path: 文件路径
        @param cor_num: 线程数量
        ''' 
        self.url = url
        self.file_path = file_path
        self.cor_num = cor_num
        self.file_size = 0
        self.precision = 100
        self.progress = None
        self.client = httpx.AsyncClient()


    def calc_split_intervals(self):
        step = self.file_size // (self.cor_num * self.precision)
        arr = list(range(0, self.file_size, step))
        arr[-1] = self.file_size
        return [ [arr[i], arr[i+1]-1] for i in range(len(arr) - 1) ]


    def get_file_size(self):
        res = httpx.head(self.url)
        self.file_size = int(res.headers['Content-Length'])


    # 借助协程异步下载
    async def async_range_download(self, s_pos, e_pos):
        headers = {"Range": f"bytes={s_pos}-{e_pos}"}
        # 下载失败的话，重复
        for i in range(0, 10):
            try:
                res = await self.client.get(self.url, headers=headers)
                self.progress.update(e_pos - s_pos)
                with open(self.file_path, "rb+") as f:
                    f.seek(s_pos)
                    f.write(res.content)
                return
            except Exception as e:
                print("请求超时，尝试重连: range:[%d, %d], retry = %d" %(s_pos, e_pos, i))
                print("报错信息：{}".format(e))
                continue
        print("下载失败：range:[%d, %d]" %(s_pos, e_pos))
        


    def start(self):
        # 先创建空文件
        with open(self.file_path, "wb") as f:
            pass

        self.get_file_size()
        self.progress = tqdm(total = self.file_size, unit_scale = True, desc = "image-or-sdk", unit_divisor = 1024, unit = 'B')
        split_intervals = self.calc_split_intervals()
        loop = asyncio.get_event_loop()
        for i in range(self.precision):
            tasks = []
            for j in range(self.cor_num):
                s_pos = split_intervals[self.cor_num * i + j][0]
                e_pos = split_intervals[self.cor_num * i + j][1]
                tasks.append(self.async_range_download(s_pos, e_pos))
            # 等待所有协程执行完毕
            loop.run_until_complete(asyncio.wait(tasks))


if __name__ == "__main__":
    # self.progress = tqdm(total = self.file_size, unit_scale = True, desc = "image-or-sdk", unit_divisor = 1024, unit = 'B')
    # self.progress.update(e_pos - s_pos)
    url = "http://download.ci.openharmony.cn/version/Daily_Version/ohos-sdk/20230322_020617/version-Daily_Version-ohos-sdk-20230322_020617-ohos-sdk.tar.gz"
    file_path = "version-Daily_Version-ohos-sdk-20230322_020617-ohos-sdk.txt"
    print("Start download the image or sdk, current time is: ", datetime.datetime.now())
    downloader = MultiCoroutineDownloader(url, file_path, 20)
    downloader.start()
    print("End download the image or sdk, current time is: ", datetime.datetime.now())
